import hashlib
import re
import time
import binascii
import uuid
import base64
import requests,json
from src import cache
from src import config 
from src.log import *
import json
import threading,traceback
import time,threading
from concurrent import futures
waiting_message={} #活跃状态
waiting_lock = threading.Lock()

gExecutor = futures.ThreadPoolExecutor(max_workers=6)


DiscordPrefix   = "<<<!"
DiscordNextFix  = "!>>>"
_global_timer = None
timer_interval = 5
def process_message_state_timer(state):
    state.process_timer()
def callback():
    global timer_interval
    Logger.info(f'Timer expired callback thread timer_interval:{timer_interval}')
    hasData = True
    with waiting_lock:
        try:
            keys_to_remove =[]
            for id, state in waiting_message.items():
                if state and state.hasData():
                    process_message_state_timer(state)
                    #gExecutor.submit(process_message_state_timer, state)
                else:
                    keys_to_remove.append(id)
            for k in keys_to_remove:
                waiting_message.pop(k, None)
            hasData = len(waiting_message) > 0
        except:
            traceback.print_exc()
        
    if hasData:    
        timer_interval = 5
    else:
        timer_interval = 15
    
    global _global_timer
    _global_timer = threading.Timer(timer_interval, callback)
    _global_timer.start()
    _global_timer.join()
            
        


    

def add_waiting_message(message):
    with waiting_lock:
        waiting_message[message.discordMsgId] = message
    global timer_interval
    timer_interval = 5




def remove_waiting_message(message):
    global timer_interval
    with waiting_lock:
         waiting_message.pop(message.discordMsgId, None)
         if len(waiting_message) <= 0:
            timer_interval = 10

def get_waiting_message(discordMsgId):
    msg = None
    with waiting_lock:
        msg = waiting_message.get(discordMsgId, None)
    return msg

def getDiscordLardMapId(content):
    re_str = r'<<<!([^!]+)!>>>'
    re_obj = re.compile(re_str)
    match = re_obj.search(content)
    if match:
        id = match.group(1)
        return id, None
    return None, Exception('not found')



    
class MessageActionState:
    def __init__(self, discordMsgId):
        self.msgId = None #对应task看 parentMsgId
        self.discordMsgId=discordMsgId #对应task的 discordMsgId
        self.msgHash =None #对应task的 msgHash
        self.currentJob=None
        self.tasks=[]
        self.UActionFlags=''
        self.lockobj = threading.Lock()

    def hasData(self):
        tasks = self.tasks
        return self.currentJob != None or (tasks is not None and len(tasks) > 0)
            

    def add_task(self,task):
        with self.lockobj:
            self.tasks.append(task)
    
    def printStr(self):
        return f'msgId={self.msgId} discordMsgId={self.discordMsgId}  msgHash={self.msgHash} currentJob={self.currentJob} UActionFlags={self.UActionFlags}'
    
    def doSendDiscordMessageMJUpscale(self, task):

        msgId = task.get('msgId', '')
        msgHash = task.get('msgHash', None)
        discordMsgId = task.get('discordMsgId', '')
        actionFlags = self.UActionFlags
        index = task.get('index',0)
        if f'U{index}' in actionFlags:
            self.currentJob = None
            return True
        error = requestMidjourney(dict(type="upscale",
        discordMsgId=discordMsgId,
        msgHash=      msgHash,
        index=        index))
        if error is None:
            self.currentJob = task
            self.UActionFlags = actionFlags + f'U{index}'
            task['timestamp'] = int(time.time())
            return True 
        else:
            self.currentJob = None
            parentMsgId = task.get('parentMsgId','')
            DiscordMsg_Callback_Failed_0(msgId, parentMsgId, msg='requestMidjourney{upscale} request error')
            return False 
        
    def doSendDiscordMessageMJVariate(self, task):
        

        msgId = task.get('msgId', '')
        msgHash = task.get('msgHash', None)
        discordMsgId = task.get('discordMsgId', '')
        index = task.get('index', 0)
        
        error = requestMidjourney(dict(type="variation",
        discordMsgId=discordMsgId,
        msgHash=      msgHash,
        index=        index))
        
        if error is None:
            self.currentJob = task
            task['timestamp'] = int(time.time())
            return True
        else:
            self.currentJob = None
            parentMsgId = task.get('parentMsgId','')
            DiscordMsg_Callback_Failed_0(msgId, parentMsgId, msg='requestMidjourney{variation} request error')
            return False    
    def doSendDiscordMessageMJMaxUpscale(self, task):
        msgId = task.get('msgId', '')
        discordMsgId = task.get('discordMsgId', '')
        msgHash = task.get('msgHash', None)
        error = requestMidjourney(dict(type="maxUpscale",
        discordMsgId=discordMsgId,
        msgHash=      msgHash))
        if error is None:
            self.currentJob = task
            task['timestamp'] = int(time.time())
            return True 
        else:
            parentMsgId = task.get('parentMsgId','')
            self.currentJob = None
            DiscordMsg_Callback_Failed_0(msgId=msgId, parentMsgId=parentMsgId, msg='requestMidjourney{MaxUpscale} request error')
            return False   
    def doSendDiscordMessageMJReset(self, task):
        
        msgId = task.get('msgId', '')
        discordMsgId = task.get('discordMsgId', None)
        msgHash = task.get('msgHash', None)
        error = requestMidjourney(dict(type="reset",
        discordMsgId=discordMsgId,
        msgHash=      msgHash))
        if error is None:
            self.currentJob = task
            task['timestamp'] = int(time.time())
            return True 
        else:
            parentMsgId = task.get('parentMsgId','')
            self.currentJob = None
            DiscordMsg_Callback_Failed_0(msgId=msgId, parentMsgId=parentMsgId, msg='requestMidjourney{reset} request error')
            return False
        
       

    @classmethod
    def create(cls, discordMsgId, dict_data):
        state = MessageActionState(discordMsgId)
        if dict_data is not None:
            state.msgId = dict_data.get('msgId', '') #对应task看 parentMsgId
            state.msgHash =dict_data.get('msgHash', '') #对应task的 msgHash
            state.currentJob=dict_data.get('currentJob', None)
            state.tasks=dict_data.get('tasks', [])
            state.UActionFlags=dict_data.get('UActionFlags', '')
        return state
    
    @classmethod
    def save(cls, state):
        if state is None:
            return
        dd = dict(
            msgId=state.msgId,
            discordMsgId=state.discordMsgId,
         msgHash=state.msgHash, 
         currentJob=state.currentJob,
         tasks=state.tasks,
          UActionFlag=state.UActionFlags,
         )
        cache.set_discord_msg_action_data(state.discordMsgId, dd)

    def process_timer(self):
        hasData = False
        t = int(time.time())
        with self.lockobj:
            currentJob = self.currentJob
            if currentJob is not None:
                expired_interval = 1.5 * 60 if currentJob.get('type', '') == 'upscale' else 15 * 60 #u执行时间比较短,v时间比较长
                if t < currentJob.get('timestamp', 0) + expired_interval: #
                    return True
                else:
                    postMJResult(dict(code=1, msgId=currentJob.get('msgId', ''), parentMsgId= currentJob.get('parentMsgId', ''), msg='cancel by timeout'))
                    self.currentJob = None
            else:
                Logger.info(f"msgId: {self.msgId}currentJob == null")
            task = None
            if self.tasks is not None and len(self.tasks) > 0:
                task = self.tasks.pop(0)
            if task is not None:
                type = task.get('type', '')
                if type == 'upscale':
                    self.doSendDiscordMessageMJUpscale(task)
                elif type == 'variation':
                    self.doSendDiscordMessageMJVariate(task)
                elif type == 'maxUpscale':
                    self.doSendDiscordMessageMJMaxUpscale(task)
                elif type == 'reset':
                    self.doSendDiscordMessageMJReset(task)
                else:
                    pass
            MessageActionState.save(self)
            hasData = self.currentJob != None or (self.tasks is not None and len(self.tasks) > 0)
        return hasData
        




def is_file_path(path):
    # 解析URL
    parsed = urllib.parse.urlparse(path)
    # 如果scheme为空，说明是文件路径
    if not parsed.scheme:
        return os.path.isfile(path)
    # 否则是URL
    return False


import imghdr
def get_img_extension_from_bytes(data):
    extension = imghdr.what(None, h=data)
    ext =  '.' + extension if extension else None
    if ext == '.jpeg':
        return '.jpg'
    return ext

 
#使用本地文件
def doMidjourneyUpload(imgData):
    size=len(imgData)
    ext = get_img_extension_from_bytes(imgData)
    if ext is not None and imgData is not None and size > 0:
        
        image_file_name = 'up_'+ str(uuid.uuid4()) + ext
        data = json.dumps(dict(name=image_file_name, size=size, imgData=base64.b64encode(imgData).decode('utf-8')))
        if ext =='.jpg':
            headers = {"Content-Type": "image/jpeg"}
        elif ext == '.png':
            headers = {"Content-Type": "image/png"}
            
        response = requests.post(config.DISCORD_MIDJOURNEY_UPLOAD_URL, headers=headers, data=data)
        json_data = response.json()
        attachment_file = json_data['name']
        return attachment_file is not None, attachment_file

    Logger.info(f'doMidjourneyUpload ext：False')
    return False, '文件不符合要求(支持jpg/png,最大不超过5m)！'




def requestMidjourney(params):
    try:
        data = json.dumps(params)
        headers = {"Content-Type": "application/json"}
        response = requests.post(config.DISCORD_MIDJOURNEY_URL, headers=headers, data=data)
        response.raise_for_status()
    except requests.exceptions.HTTPError as e:
        print("HTTP Error:", e)
        Logger.error(str(e), exc_info=True)
        return e
    except requests.exceptions.ConnectionError as e:
        print("Error Connecting:", e)
        Logger.error(str(e), exc_info=True)
        return e
    except requests.exceptions.Timeout as e:
        print("Timeout Error:", e)
        Logger.error(str(e), exc_info=True)
        return e
    except requests.exceptions.RequestException as e:
        print("Something went wrong:", e)
        Logger.error(str(e), exc_info=True)
        return e
    except Exception as e:
        return e
    return None

def SendDiscordMessageMJGenerate(conversation_id, msgId, parentMsgId, prompt):
    str1 = msgId + str(time.time_ns())
    hash = hashlib.md5(str1.encode())
    id = binascii.hexlify(hash.digest())[:12].decode()
    prompt = DiscordPrefix + id + DiscordNextFix + prompt
    error = requestMidjourney(dict(type='generate', prompt=prompt))
    if error is None:
        data = dict()
        data['msgId'] = msgId
        data['parentMsgId'] = parentMsgId #应该为空或None
        data['conversation_id'] = conversation_id
        cache.set_discord_image_session(id, data)
        return id, dict(code=0,msg='')
    else:
        return id, dict(code=1, msg='MJ任务提交失败。请仔细检查指令和参数是否正常！')

    
def get_message(discordMsgId): 
    message_action_data_state = get_waiting_message(discordMsgId)
    if message_action_data_state is None:
        cache_data = cache.get_discord_msg_action_data(discordMsgId)
        message_action_data_state = MessageActionState.create(discordMsgId, cache_data)
    return message_action_data_state
    

    
def SendDiscordMessageMJUpscale(msgId, parentMsgId, discordMsgId, msgHash, index):
    task = dict(isFromU = True, timestamp=0, type='upscale', msgId=msgId, parentMsgId=parentMsgId, discordMsgId=discordMsgId, msgHash=msgHash, index=index)
    message_action_data_state = get_message(discordMsgId)
    message_action_data_state.discordMsgId = discordMsgId
    message_action_data_state.msgId = msgId
    message_action_data_state.msgHash = msgHash
    currentJob = message_action_data_state.currentJob
    message_action_data_state.add_task(task)
    add_waiting_message(message_action_data_state)
    return dict(code=0,msg='')
        

#msgId 触发用户动作所在消息的ID.
def SendDiscordMessageMJVariate(msgId, parentMsgId, discordMsgId, msgHash,  index):
    
    task = dict(isFromU = False, timestamp=0, type='variation', msgId=msgId, parentMsgId=parentMsgId, discordMsgId=discordMsgId, msgHash=msgHash, index=index)
    message_action_data_state = get_message(discordMsgId)
    message_action_data_state.discordMsgId = discordMsgId
    message_action_data_state.msgId = msgId
    message_action_data_state.msgHash = msgHash
    currentJob = message_action_data_state.currentJob
    message_action_data_state.add_task(task)
    add_waiting_message(message_action_data_state)
    return dict(code=0,msg='')

def SendDiscordMessageMJMaxUpscale(msgId, parentMsgId,  discordMsgId, msgHash):
    task = dict(isFromU = False, timestamp=0, type='maxUpscale', msgId=msgId, parentMsgId=parentMsgId, discordMsgId=discordMsgId, msgHash=msgHash, index=0)
    message_action_data_state = get_message(discordMsgId)
    message_action_data_state.discordMsgId = discordMsgId
    message_action_data_state.msgId = msgId
    message_action_data_state.msgHash = msgHash
    currentJob = message_action_data_state.currentJob
    message_action_data_state.add_task(task)
    add_waiting_message(message_action_data_state)
    return dict(code=0,msg='')
      
    

     

#msgId 触发用户动作所在消息的ID.
def SendDiscordMessageMJReset(msgId, parentMsgId, discordMsgId, msgHash):
    
    task = dict(isFromU = False, timestamp=0, type='reset', msgId=msgId, parentMsgId=parentMsgId, discordMsgId=discordMsgId, msgHash=msgHash, index=0)
    message_action_data_state = get_message(discordMsgId)
    message_action_data_state.discordMsgId = discordMsgId
    message_action_data_state.msgId = msgId
    message_action_data_state.msgHash = msgHash
    currentJob = message_action_data_state.currentJob
    message_action_data_state.add_task(task)
    add_waiting_message(message_action_data_state)
    return dict(code=0,msg='')



def doSendDiscordMessageDescribe(imageData):
    id = None
    try:
        isOk, attachments_file = doMidjourneyUpload(imageData)
        if isOk:
            id = attachments_file[attachments_file.rfind('/')+1:]
            params = dict(type='describe', prompt=attachments_file)
            data = json.dumps(params)
            headers = {"Content-Type": "application/json"}
            response = requests.post(config.DISCORD_MIDJOURNEY_URL, headers=headers, data=data)
            response.raise_for_status()
        else:
            return Exception('文件数据上传失败'), None
    except requests.exceptions.HTTPError as e:
        print("HTTP Error:", e)
        Logger.error(str(e), exc_info=True)
        return e, None
    except requests.exceptions.ConnectionError as e:
        print("Error Connecting:", e)
        Logger.error(str(e), exc_info=True)
        return e, None
    except requests.exceptions.Timeout as e:
        print("Timeout Error:", e)
        Logger.error(str(e), exc_info=True)
        return e, None
    except requests.exceptions.RequestException as e:
        print("Something went wrong:", e)
        Logger.error(str(e), exc_info=True)
        return e, None
    except Exception as e:
        return e, None
    return None, id
    
def download_image(imageUrl):
    response = requests.get(imageUrl, stream=True)
    if response.status_code == 200:
        content = bytearray()
        for chunk in response.iter_content(chunk_size=1024):
            content += chunk
        return content
    return None

def SendDiscordMessageDescribe(conversation_id, msgId, imageUrl, imageData):
    if imageData is None or len(imageData) <= 0:
        imageData = download_image(imageUrl) 
    
    err, id  = doSendDiscordMessageDescribe(imageData)
    if err is None:
        job_session =dict(id=id, conversation_id=conversation_id, imageUrl=imageUrl, msgId=msgId)
        cache.set_discord_image_session(id, job_session)
    else:
        DiscordMsg_Callback_Failed_0(msgId, '', 'MidJourney指令处理出错！请稍后重试！')




def DiscordMsg_Callback_OnDescribeSuccess(id, description):
    job_session  = cache.get_discord_image_session(id)
    if job_session:
        
        msgId = job_session.get('msgId', None)
        imageUrl = job_session.get('imageUrl', None)
        conversation_id = job_session.get('conversation_id', None)
        data = dict(type='describe', msgId=msgId, imageUrl=imageUrl, conversation_id=conversation_id, description=description)
        postMJResult(data)
        return dict(code=0, msg='')
    else:
        Logger.info(f'DiscordMsg_Callback_OnDescribeSuccess id:{id} == null')
    
def DiscordMsg_Callback_OnImageEnd(id, thumailImageUrl, fullImageUrl, discordMsgId, msgHash, parent_discord_msgId):
    try:
        isFromU = False
        msgId = None
        parentMsgId = None
        #首次生成
        if parent_discord_msgId is None: # #imagen命令
            session_data = cache.get_discord_image_session(id)
            msgId = session_data.get('msgId', None)
            message_action_data_state = get_message(discordMsgId) #应该为空
            with message_action_data_state.lockobj:
                message_action_data_state.tasks=[]
                message_action_data_state.UActionFlags=''
                message_action_data_state.currentJob = None
                parentMsgId = None    
                MessageActionState.save(message_action_data_state)
            result = DiscordMsg_Callback_MaxUpscale(msgId, parentMsgId, thumailImageUrl, fullImageUrl, discordMsgId, msgHash, parent_discord_msgId) 
            return result
        else:
            message_action_data_state = get_message(parent_discord_msgId) #非imagen指令情形下-才生效
            currentJob = message_action_data_state.currentJob
            if currentJob is not None:
                index = currentJob.get('index', 0)
                isFromU = currentJob.get('isFromU',False)
                msgId = currentJob.get('msgId','')
                parentMsgId = currentJob.get('parentMsgId', '') #== message_action_data_state.msgId
                Logger.info(f'CURRENT JOB process msgId:{msgId} success ')
                if isFromU:
                    result = DiscordMsg_Callback_Upscale(msgId, parentMsgId, thumailImageUrl, fullImageUrl, discordMsgId, msgHash, index, parent_discord_msgId)
                else:
                    result = DiscordMsg_Callback_MaxUpscale(msgId, parentMsgId, thumailImageUrl, fullImageUrl, discordMsgId, msgHash, parent_discord_msgId)
                
                with message_action_data_state.lockobj:
                    message_action_data_state.currentJob = None
                    MessageActionState.save(message_action_data_state)
                return result
            else:
                with message_action_data_state.lockobj:
                    message_action_data_state.currentJob = None
                    MessageActionState.save(message_action_data_state)
                Logger.info(f'DiscordMsg_Callback_OnImageEnd {id}: currentJob== null')
                return dict(code=1, msg='应该是超时导致的相应的Job被取消了。在取消时，已经发过回调了。所以不需要再重新发。')
    except Exception as e:
        Logger.error(str(e), exc_info=True)
        traceback.print_exc()
    return dict(code=1, msg='出错') 
    

def DiscordMsg_Callback_Failed_0(msgId, parentMsgId, msg=''):
     Logger.info(f'DiscordMsg_Callback_Failed_0 msgId:{msgId} parentMsgId:{parentMsgId} msg：{msg}')
     postMJResult(dict(code=1, msgId=msgId, parentMsgId= '' if parentMsgId is None else parentMsgId, msg=msg))
     return dict(code=0, msg='🤖️ ：抱歉,Midjourney任务执行出错。请稍后重试') 

def DiscordMsg_Callback_Failed(id, discordMsgId, parent_discord_msgId):
    msgId = None
    parentMsgId = None
    if parent_discord_msgId is not None: #理论上还是可能存在对应不上的问题，只能忽略吧。
        message_action_data_state = get_message(parent_discord_msgId)
        parentMsgId = message_action_data_state.msgId
        currentJob =  message_action_data_state.currentJob
        if currentJob is not None: 
            msgId = currentJob.get('msgId', None)
            parentMsgId = currentJob.get('parentMsgId', None) 
    else:
        session_data = cache.get_discord_image_session(id)
        msgId = session_data.get('msgId', None)
        parentMsgId = session_data.get('parentMsgId', None) 
    
    with message_action_data_state.lockobj:
        message_action_data_state.currentJob = None 
        MessageActionState.save(message_action_data_state)
    result = DiscordMsg_Callback_Failed_0(msgId, parentMsgId=parentMsgId, msg='🤖️ ：抱歉,Midjourney任务执行出错。请稍后重试')
    
    return result
    

def postMJResult(data):
    try:
        headers = {"Content-Type": "application/json"}
        msgId = data['msgId'] if 'msgId' in data else None
        if msgId is not None and msgId.startswith('wx__'):
            response = requests.post(config.CLIENT_WX_CALLBACK_URL, headers=headers, data=json.dumps(data))
            
        else:
            response = requests.post(config.CLIENT_CALLBACK_URL, headers=headers, data=json.dumps(data))
            
        response.raise_for_status()
    except requests.exceptions.HTTPError as errh:
        print("HTTP Error:", errh)
        return errh
    except requests.exceptions.ConnectionError as errc:
        print("Error Connecting:", errc)
        return errc
    except requests.exceptions.Timeout as errt:
        print("Timeout Error:", errt)
        return errt
    except requests.exceptions.RequestException as err:
        print("Something went wrong:", err)
        return err
    return None

#图片地址：消息
def DiscordMsg_Callback_MaxUpscale(msgId, parentMsgId, imageUrl, fullImageUrl, discordMsgId, msgHash, parent_discord_msgId):
    err = postMJResult(dict(type='image', code=0, msgId=msgId, parentMsgId='' if parentMsgId is None else parentMsgId, imageUrl=imageUrl, fullImageUrl=fullImageUrl, discordMsgId=discordMsgId, msgHash=msgHash,upscale=True, index=0, parent_discord_msgId=parent_discord_msgId))
    if err is None:
        return dict(code=0, msg='')
    return dict(code=1, msg='')

def DiscordMsg_Callback_Upscale(msgId, parentMsgId, imageUrl, fullImageUrl, discordMsgId, msgHash, index,parent_discord_msgId):
    err = postMJResult(dict(type='image', code=0, msgId=msgId, parentMsgId='' if parentMsgId is None else parentMsgId, imageUrl=imageUrl, fullImageUrl=fullImageUrl, discordMsgId=discordMsgId, msgHash=msgHash,upscale=False, index=index,parent_discord_msgId=parent_discord_msgId))
    if err is None:
        return dict(code=0, msg='')
    return dict(code=1, msg='')


def init_timer_work():
    global _global_timer
    try:
        if _global_timer:
            _global_timer.cancel()
            _global_timer = None
    except:
        traceback.print_exc()

    def timer_work():
        _global_timer = threading.Timer(5, callback)
        _global_timer.start()
        _global_timer.join()
    gExecutor.submit(timer_work)
    

